import math
import numpy as np
from matplotlib import pyplot

def cdatos2(datos):
    datos2=[]
    for dt in datos:
        datos2.append([dt[0],dt[1]*2.0,dt[2]*2.0,dt[3]*2.0])
    return datos2

def anint(x, lx):
    res=x
    if x>0.5*lx:
        res-=lx
    elif x<-0.5*lx:
        res+=lx
    return res

def distancia(lbox,datos,i,j):
    dis2=0
    drx=datos[i][1]-datos[j][1]
    drx=anint(drx, lbox)
    dry=datos[i][2]-datos[j][2]
    dry=anint(dry, lbox)
    drz=datos[i][3]-datos[j][3]
    drz=anint(drz, lbox)
    dis2=drx*drx+dry*dry+drz*drz
    return  [math.sqrt(dis2),drx,dry,drz]

def load(name):
    f = open(name,'r')
    tx = f.readline()
    tx = tx.replace('\n','')
    tx = tx.replace('{','')
    tx = tx.replace('}','')
    val = tx.split(' ')
    aux=[]
    for v in val:
        if v!='':
            aux.append(v)
    npart=int(aux[1])
    lbox=float(aux[3])
    eTotal=float(aux[5])
    ekin=float(aux[7])
    emag=float(aux[9])
    erep=float(aux[11])
    
    line = f.readline()
    datos=[]
    for i in range(npart):
        line = f.readline()
        line=line.replace('\n','')
        line=line.replace('{','')
        line=line.replace('}','')
        val=line.split(' ')
        res=[]
        for v in val:
            if v!='':
                res.append(float(v))
        datos.append(res)
    
    f.close()
    return [npart, lbox, eTotal, ekin, emag, erep,  datos]

def eKin(datos):
    ek=0.0;
    np=0 
    for dt in datos:
        np+=1
        vx=dt[4]
        vy=dt[5]
        vz=dt[6]
        v2=vx*vx+vy*vy+vz*vz
        ek+=v2
    return 0.5*ek/len(datos)

def force2(rcut, lbox, datos):
    forc=np.zeros((len(datos),3))  
    for i in range(len(datos)):
        for j in range(len(datos)):
            if i!=j:
                [dis,dx,dy,dz] = distancia(lbox,datos,i,j)
                if dis<rcut:
                    f=24.*(2.*math.pow(dis,-13.0)-math.pow(dis,-7.0))
                    fx=f*dx/dis
                    fy=f*dy/dis
                    fz=f*dz/dis        
                
                    forc[i][0]+=fx
                    forc[i][1]+=fy
                    forc[i][2]+=fz
    return forc

def forceDipij(m,lbox, datos,i,j):
    u0=m*m
    [dis,dx,dy,dz] = distancia(lbox,datos,i,j)
    Co= dz/dis
    Co2=Co*Co
    fx=-3.*u0*(5.*Co2-1.)*dx/math.pow(dis,5.0)
    fy=-3.*u0*(5.*Co2-1.)*dy/math.pow(dis,5.0)
    fz=-3.*u0*(5.*Co2-3.)*dz/math.pow(dis,5.0)
    return [fx,fy,fz]

#def forceDip(m,lbox, datos):
#    u0=m*m
#    forc=np.zeros((len(datos),3))  
#    for i in range(len(datos)):
#        for j in range(len(datos)):
#            if i!=j:
#                [dis,dx,dy,dz] = distancia(lbox,datos,i,j)
#                Co= dz/dis
#                Co2=Co*Co
#                fx=-3.*u0*(5.*Co2-1.)*dx/math.pow(dis,5.0)
#                fy=-3.*u0*(5.*Co2-1.)*dy/math.pow(dis,5.0)
#                fz=-3.*u0*(5.*Co2-3.)*dz/math.pow(dis,5.0)
#            
#                forc[i][0]+=fx
#                forc[i][1]+=fy
#                forc[i][2]+=fz
#    return forc

def forceDip(m,lbox, datos):
    u0=m*m
    forc=np.zeros((len(datos),3))  
    for i in range(len(datos)):
        for j in range(len(datos)):
            if i!=j:
                [fx,fy,fz]=forceDipij(m,lbox, datos,i,j)            
                forc[i][0]+=fx
                forc[i][1]+=fy
                forc[i][2]+=fz
    return forc

def stressTensor(m,lbox,datos,volume):
    npart = len(datos)
    sumv=np.zeros((3,3))
    sumf=np.zeros((3,3))
    sigma=np.zeros((3,3))
    for i in range(npart):
        vi=datos[i][3:6:1]
        for i1 in range(3):
            for i2 in range(3):
                sumv[i1][i2]+=vi[i1]*vi[i2]
    for i in range(npart-1):
        for j in range(i+1,npart):
            force=forceDipij(m,lbox, datos,i,j)
            [dis,dx,dy,dz] = distancia(lbox,datos,i,j)
            dis=[dx,dy,dz]
            for i1 in range(3):
                for i2 in range(3):
                    sumf[i1][i2]+=force[i2]*dis[i1]
                    #print dis[i1],dis[i2]
                    #print force[i1],force[i2]
    for i1 in range(3):
            for i2 in range(3):
                sigma[i1][i2]=(sumv[i1][i2]+sumf[i1][i2])/volume
    return [sumv,sumf,sigma]

def Epoten(rcut,lbox,datos):
    ePot=0.
    for i in range(len(datos)):
        for j in range(len(datos)):
            if i!=j:
                res= distancia(lbox,datos,i,j)
                dis=res[0]     
                potdm=0.           
                if dis<rcut:
                    dis3=dis*dis*dis
                    dis6=dis3*dis3
                    dis12=dis6*dis6
                    potdm = (1./dis12-1./dis6)
                    #print i,j,dis,potdm
                ePot+=potdm
    return 2.*ePot

def EpotenDip(m,lbox,datos):
    u0=m*m
    ePot=0.
    for i in range(len(datos)):
        for j in range(len(datos)):
            if i!=j:
                [dis, dx,dy,dz]= distancia(lbox,datos,i,j)
                Co= dz/dis
                Co2=Co*Co
                dis3=dis*dis*dis
                potdm = u0*(1.-3.*Co2)/dis3
                    #print i,j,dis,potdm
                ePot+=potdm
    return 0.5*ePot#/len(datos)




def testEnergy(m,lbox,datos):
    xnew=[]
    ynew=[]
    nc=len(datos)
    x=[]
    y=[]
    n=0
    for dt in datos:
        x.append(n)
        y.append(dt[1])
        n+=1
    
    for i in range(nc):
        name='/home/jonk/codeEspresso/test/posvelforce'+str(i)+'.dat'
        datos = load(name)
        Ek=eKin(datos)
        Temp=Ek*2./(3.*300)
        rcut=2.0
        ePot=EpotenDip(m,lbox,datos)
        xnew.append(i)
        ynew.append(ePot)
        #print i,Ek, ePot
    
    
    
    pyplot.plot(x,y,'ro')
    pyplot.plot(xnew,ynew,'bx')
    pyplot.show()

def testForceMagnetic(m,lbox, name):
    datos = load(name)
    forcesEspresso=[]
    for dt in datos:
        forcesEspresso.append([dt[6],dt[7],dt[8]])
    ##print forcesEspresso[0]
    forceMio = forceDip(m,lbox,datos)
    ##print forceMio[0]
    x=[]
    y1=[]
    y2=[]
    for i in range(len(forceMio)):
        x.append(i)
        y1.append(forcesEspresso[i][2])
        y2.append(forceMio[i][2])
    #print y1
    #print y2
    pyplot.plot(x,y1,'ro')
    pyplot.plot(x,y2,'bx')
    pyplot.show()


#Para comparar con las fuerzas de "Espresso" es necesario poner el dumping del termostato a 0
#para obtener unicamente las fuerzas debidas al potencial

u0=20.0/4.0
rcut = math.pow(2,1./6.)
m=math.sqrt(u0)
name='/home/jonk/eclipse/TCLMagnetic/config_0.dat'
#testForceMagnetic(m,lbox, name)
[npart, lbox, eTotal, ekin, emag, erep, datos] = load(name)
ekinOld = eKin(datos)
eDipOld = EpotenDip(m,lbox, datos)
epoten = Epoten(rcut,lbox, datos)
print 'eKinOld: ',ekinOld,' EkinEspresso: ',ekin
print 'eDipOld: ',eDipOld,' EDipEspresso: ',emag 
#print 'eRepOld: ',epoten,' ERepEspresso: ',erep
#[sumv,sumf,sigma] = stressTensor(m,lbox,datos,volume)
u02=u0*8.0
m=math.sqrt(u02)
datos2=cdatos2(datos)
xxx = EpotenDip(m,2.*lbox, datos2)
print 'Emagnetica JP: ',xxx

m=math.sqrt(u0)    
uno =forceDipij(m,lbox, datos,0,5)
print 'Fuerza espresso: ',uno
m=math.sqrt(u02)
dos=forceDipij(m,2.*lbox, datos2,0,5)
print 'Fuerza JP: ',dos
print uno[0]/dos[0],uno[1]/dos[1],uno[2]/dos[2]
 


